from tqdm import tqdm

entities = [str(1)]

with open('facts.txt', 'r') as f:
    x = f.readlines()

ent_toy = list()
with open('facts_toy.txt', 'w') as ff:
    for l in tqdm(x):
        h, r, t = l.split()
        if h.strip() in entities or t.strip() in entities:
            ff.write('{}'.format(l))
            print(h.strip(), t.strip())
            ent_toy.append(h.strip())
            ent_toy.append(t.strip())

ent_toy = list(set(ent_toy))
with open('entities_toy.txt', 'w') as fff:
    for e in ent_toy:
        fff.write('{}\n'.format(e))


